#include "Fluid2D.h"

Fluid2D::Fluid2D(void)
{
}

Fluid2D::~Fluid2D(void)
{
}

void Fluid2D::Setup( int n, float dt )
{
	this->n = n;
	this->dt = dt;       

	density=(float**)malloc(n*sizeof(float*));
	pre_density=(float**)malloc(n*sizeof(float*));
	curl=(float**)malloc(n*sizeof(float*));
	u=(float**)malloc(n*sizeof(float*));
	pre_u=(float**)malloc(n*sizeof(float*));
	v=(float**)malloc((n+1)*sizeof(float*));
	pre_v=(float**)malloc((n+1)*sizeof(float*));
	A=(float**)malloc(n*n*sizeof(float*));
	entry=(int**)malloc(n*n*sizeof(int*));
	b=(float*)malloc(n*n*sizeof(float));
	buffer=(float*)malloc(n*n*sizeof(float));
	temp=(float*)malloc(n*n*sizeof(float));

	for (int i=0;i<n;i++)
	{
		density[i]=(float*)malloc(n*sizeof(float));
		pre_density[i]=(float*)malloc(n*sizeof(float));
		curl[i]=(float*)malloc(n*sizeof(float));
		u[i]=(float*)malloc((n+1)*sizeof(float));
		pre_u[i]=(float*)malloc((n+1)*sizeof(float));
		v[i]=(float*)malloc(n*sizeof(float));
		pre_v[i]=(float*)malloc(n*sizeof(float));
	}
	v[n]=(float*)malloc(n*sizeof(float));
	pre_v[n]=(float*)malloc(n*sizeof(float));

	for (int i=0;i<n*n;i++)
	{
		A[i]=(float*)malloc(5*sizeof(float));
		entry[i]=(int*)malloc(5*sizeof(int));

		for (int j=0;j<5;j++)
		{
			A[i][j]=entry[i][j]=0;
		}
		b[i]=buffer[i]=temp[i]=0;
		
	}


	int index;

	for(int i=0;i<n;i++)
	{        	
		for(int j=0;j<n;j++)
		{
			u[i][j]=pre_u[i][j]=v[i][j]=pre_v[i][j]=density[i][j]=pre_density[i][j]=0;
			curl[i][j]=0;

			index=i*n+j;

			A[index][0]=0;

			if(i-1>=0){
				A[index][0]++;
				A[index][1]=-1;
				entry[index][1]=(i-1)*n+j;
			}
			if(i+1<n){
				A[index][0]++;
				A[index][2]=-1;
				entry[index][2]=(i+1)*n+j;
			}
			if(j-1>=0){
				A[index][0]++;
				A[index][3]=-1;
				entry[index][3]=i*n+j-1;
			}
			if(j+1<n){
				A[index][0]++;
				A[index][4]=-1;
				entry[index][4]=i*n+j+1;
			}


		}
		u[i][n]=pre_u[i][n]=v[n][i]=pre_v[n][i]=0;
	}

}

void Fluid2D::ApplyBuoyancy( float** Fbuoy )
{
	float Tamb = 0;
	float a = 0.000625f;
	float b = 0.025f;

	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			Tamb += density[i][j];
		}
	}

	Tamb /= (n * n);

	for (int i = 1; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			float buoy=0.5f*(density[i-1][j]+density[i][j]);
			Fbuoy[i][j]=a*buoy-b*(buoy-Tamb);
		}
	}
}

void Fluid2D::AddSourceU( float ** x, float ** x0 )
{
	for (int i = 0; i < n; i++)
	{
		for(int j=1;j<n;j++)
		{
			x[i][j]+=dt*x0[i][j];
		}
	}
}

void Fluid2D::AddSourceV( float ** x, float ** x0 )
{
	for (int i = 1; i < n; i++)
	{
		for(int j=0;j<n;j++)
		{
			float deb=x0[i][j];
			float debb=x[13][0];
			x[i][j]+=dt*x0[i][j];
		}
	}
}

void Fluid2D::AddSourceD( float ** x, float ** x0 )
{
	for (int i = 0; i < n; i++)
	{
		for(int j=0;j<n;j++)
		{
			x[i][j]+=dt*x0[i][j];
		}
	}
}

void Fluid2D::AdvectU( int b, float** d, float** d0, float** du, float** dv )
{
	int i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	dt0 = dt * n;

	for(int i=0;i<n;i++)
	{
		d[i][0]=d[i][n]=0;
	}

	for (int i = 0; i < n; i++)
	{
		for (int j = 1; j < n; j++)
		{
			// go backwards through velocity field
			y = j - dt0 * du[i][j];

			float inter=(dv[i][j-1]+dv[i][j]+dv[i+1][j-1]+dv[i+1][j])*0.25f;
			x = i - dt0 * inter;

			if(x>=0&&x<n-1&&y>=0&&y<n)
			{
				i0=(int)x;
				i1=i0+1;

				j0 = (int) y;
				j1 = j0 + 1;

				s1 = x - i0;
				s0 = 1 - s1;
				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j] = s0 * (t0 * d0[i0][j0] + t1 * d0[i0][j1])
					+ s1 * (t0 * d0[i1][j0] + t1 * d0[i1][j1]);
			}else if(x<0&&y>=0&&y<n)
			{                	
				j0 = (int) y;
				j1 = j0 + 1;

				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j]=t0*d0[0][j0]+t1*d0[0][j1];

			}else if(x>=n-1&&y>=0&&y<n)
			{
				j0 = (int) y;
				j1 = j0 + 1;

				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j]=t0*d0[n-1][j0]+t1*d0[n-1][j1];
			}else if(y<0&&x>=0&&x<n-1)
			{
				i0=(int)x;
				i1=i0+1;

				s1 = x - i0;
				s0 = 1 - s1;

				d[i][j]=s0*d0[i0][0]+s1*d0[i1][0];
			}else if(y>=n&&x>=0&&x<n-1)
			{
				i0=(int)x;
				i1=i0+1;

				s1 = x - i0;
				s0 = 1 - s1;

				d[i][j]=s0*d0[i0][n]+s1*d0[i1][n];
			}else if(x<0&&y<0)
			{
				d[i][j]=d0[0][0];
			}else if(x<0&&y>=n)
			{
				d[i][j]=d0[0][n];
			}else if(x>=n-1&&y<0)
			{
				d[i][j]=d0[n-1][0];
			}else if(x>=n-1&&y>=n)
			{
				d[i][j]=d0[n-1][n];
			}

		}
	}
}

void Fluid2D::AdvectV( int b, float** d, float** d0, float** du, float** dv )
{
	int i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	dt0 = dt * n;

	for(int i=0;i<n;i++)
	{
		d[0][i]=d[n][i]=0;
	}

	for (int i = 1; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			// go backwards through velocity field
			x = i - dt0 * dv[i][j];

			float inter=(du[i-1][j+1]+du[i-1][j]+du[i][j+1]+du[i][j])*0.25f;
			y = j - dt0 * inter;

			if(x>=0&&x<n&&y>=0&&y<n-1)
			{
				i0=(int)x;
				i1=i0+1;

				j0 = (int) y;
				j1 = j0 + 1;

				s1 = x - i0;
				s0 = 1 - s1;
				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j] = s0 * (t0 * d0[i0][j0] + t1 * d0[i0][j1])
					+ s1 * (t0 * d0[i1][j0] + t1 * d0[i1][j1]);
			}else if(x<0&&y>=0&&y<n-1)
			{                	
				j0 = (int) y;
				j1 = j0 + 1;

				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j]=t0*d0[0][j0]+t1*d0[0][j1];

			}else if(x>=n&&y>=0&&y<n-1)
			{
				j0 = (int) y;
				j1 = j0 + 1;

				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j]=t0*d0[n][j0]+t1*d0[n][j1];
			}else if(y<0&&x>=0&&x<n)
			{
				i0=(int)x;
				i1=i0+1;

				s1 = x - i0;
				s0 = 1 - s1;

				d[i][j]=s0*d0[i0][0]+s1*d0[i1][0];
			}else if(y>=n-1&&x>=0&&x<n)
			{
				i0=(int)x;
				i1=i0+1;

				s1 = x - i0;
				s0 = 1 - s1;

				d[i][j]=s0*d0[i0][n-1]+s1*d0[i1][n-1];
			}else if(x<0&&y<0)
			{
				d[i][j]=d0[0][0];
			}else if(x<0&&y>=n-1)
			{
				d[i][j]=d0[0][n-1];
			}else if(x>=n&&y<0)
			{
				d[i][j]=d0[n][0];
			}else if(x>=n&&y>=n-1)
			{
				d[i][j]=d0[n][n-1];
			}

		}
	}
}

void Fluid2D::AdvectD( int b, float** d, float** d0, float** du, float** dv )
{
	int i0, j0, i1, j1;
	float x, y, s0, t0, s1, t1, dt0;

	dt0 = dt * n;

	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			// go backwards through velocity field
			float interV=(dv[i][j]+dv[i+1][j])*0.5f;
			x = i - dt0 * interV;

			float interU=(du[i][j]+du[i][j+1])*0.5f;
			y = j - dt0 * interU;

			if(x>=0&&x<n-1&&y>=0&&y<n-1)
			{
				i0=(int)x;
				i1=i0+1;

				j0 = (int) y;
				j1 = j0 + 1;

				s1 = x - i0;
				s0 = 1 - s1;
				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j] = s0 * (t0 * d0[i0][j0] + t1 * d0[i0][j1])
					+ s1 * (t0 * d0[i1][j0] + t1 * d0[i1][j1]);
			}else if(x<0&&y>=0&&y<n-1)
			{                	
				j0 = (int) y;
				j1 = j0 + 1;

				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j]=t0*d0[0][j0]+t1*d0[0][j1];

			}else if(x>=n-1&&y>=0&&y<n-1)
			{
				j0 = (int) y;
				j1 = j0 + 1;

				t1 = y - j0;
				t0 = 1 - t1;

				d[i][j]=t0*d0[n-1][j0]+t1*d0[n-1][j1];
			}else if(y<0&&x>=0&&x<n-1)
			{
				i0=(int)x;
				i1=i0+1;

				s1 = x - i0;
				s0 = 1 - s1;

				d[i][j]=s0*d0[i0][0]+s1*d0[i1][0];
			}else if(y>=n-1&&x>=0&&x<n-1)
			{
				i0=(int)x;
				i1=i0+1;

				s1 = x - i0;
				s0 = 1 - s1;

				d[i][j]=s0*d0[i0][n-1]+s1*d0[i1][n-1];
			}else if(x<0&&y<0)
			{
				d[i][j]=d0[0][0];
			}else if(x<0&&y>=n-1)
			{
				d[i][j]=d0[0][n-1];
			}else if(x>=n-1&&y<0)
			{
				d[i][j]=d0[n-1][0];
			}else if(x>=n-1&&y>=n-1)
			{
				d[i][j]=d0[n-1][n-1];
			}

		}
	}
}

void Fluid2D::Project( float** x, float** y )
{
	int index;

	for(int i=0;i<n;i++)
	{
		for(int j=0;j<n;j++)
		{
			index=i*n+j;

			b[index]=0;

			if(i-1>=0){
				b[index]-=y[i][j];
			}
			if(i+1<n){
				b[index]+=y[i+1][j];
			}
			if(j-1>=0){
				b[index]-=x[i][j];
			}
			if(j+1<n){
				b[index]+=x[i][j+1];
			}

			b[index]*=(-1/dt);
		}
	}

	int num=n*n;

	int loop;

	for (loop=0;loop<100;loop++)
	{
		for (int i=0;i<num;i++)
		{
			float alpha=0.0f;

			//alpha=A[i][1]*buffer[(i-1)*n+j]+A[i][2]*buffer[(i+1)*n+j];

			for (int j=1;j<5;j++)
			{
				alpha+=A[i][j]*buffer[entry[i][j]];
				/*if(i!=j){
				alpha+=A[i][j]*buffer[j];
				}*/
				//alpha+=jello->MA[i][j]*jello->buffer[jello->sparse[i][j]];
			}

			if (A[i][0]!=0)
			{
				temp[i]=(b[i]-alpha)/A[i][0];
			}			
		}

		float z=abs(temp[0]-buffer[0]);

		for (int i=0;i<num;i++)
		{
			float compare=abs(temp[i]-buffer[i]);

			if(z<compare)
			{
				z=compare;
			}
			buffer[i]=temp[i];
			//jello->buffer[i]=temp[i];
		}

		if (z<1e-5)
		{
			break;
		}
	}

	for(int i=0;i<n;i++)
	{
		for(int j=1;j<n;j++)
		{
			x[i][j]-=(buffer[i*n+j]-buffer[i*n+j-1])*dt;
			y[j][i]-=(buffer[j*n+i]-buffer[(j-1)*n+i])*dt;
		}
	}
}

float Fluid2D::CalCurl( int i, int j )
{
	float du_dy = (u[i][j+1]-u[i][j])*0.5f;// (u[I(i, j + 1)] - u[I(i, j - 1)]) * 0.5f;
	float dv_dx = (v[i+1][j]-v[i][j])*0.5f;// (v[I(i + 1, j)] - v[I(i - 1, j)]) * 0.5f;

	return du_dy - dv_dx;
}

void Fluid2D::VorticityConfinement( float** Fvc_x, float** Fvc_y )
{
	float dw_dx, dw_dy;
	float length;
	float v;

	// Calculate magnitude of curl(u,v) for each cell. (|w|)
	for (int i = 0; i < n; i++)
	{
		for (int j = 0; j < n; j++)
		{
			curl[i][j] = abs(CalCurl(i, j));

			if(j!=0)
			{
				Fvc_x[i][j]=0;
			}

			if(i!=0)
			{
				Fvc_y[i][j]=0;
			}
		}
	}

	for(int i=1;i<n-1;i++)
	{
		for(int j=1;j<n-1;j++)
		{
			dw_dx = (curl[i+1][j] - curl[i-1][j]) * 0.5f;
			dw_dy = (curl[i][j+1] - curl[i][j-1]) * 0.5f;

			length=sqrt(dw_dx * dw_dx + dw_dy * dw_dy)+0.000001f;

			dw_dx /= length;
			dw_dy /= length;

			v = CalCurl(i, j);

			float bufx=dw_dy*(-v)*0.5f;
			float bufy=dw_dx*v*0.5f;

			Fvc_x[i][j]+=bufx;
			Fvc_x[i][j+1]+=bufx;
			Fvc_y[i][j]+=bufy;
			Fvc_y[i+1][j]+=bufy;

			//Fvc_x[I(i, j)] = dw_dy * -v;
			//Fvc_y[I(i, j)] = dw_dx *  v;
		}
	}
}

void Fluid2D::VelocitySolver()
{
	AddSourceU(u, pre_u);
	AddSourceV(v, pre_v);

	// add in vorticity confinement force
	VorticityConfinement(pre_u, pre_v);
	AddSourceU(u, pre_u);
	AddSourceV(v, pre_v);

	// add in buoyancy force
	ApplyBuoyancy(pre_u);
	AddSourceU(u, pre_u);

	

	
	Project(u,v);


	SwapU(); SwapV();

	// self advect velocities
	AdvectU(1, u, pre_u, pre_u, pre_v);
	AdvectV(2, v, pre_v, pre_u, pre_v);

	
	Project(u,v);

	// clear all input velocities for next frame
	for(int i=0;i<n;i++)
	{
		for(int j=0;j<n;j++)
		{
			pre_u[i][j]=pre_v[i][j]=0;
		}
		pre_u[i][n]=pre_v[n][i]=0;
	}
}

void Fluid2D::DensitySolver()
{
	
	AddSourceD(density, pre_density);

	
	SwapD();

	

	AdvectD(0, density, pre_density, u, v);

	

	// clear input density array for next frame
	for(int i=0;i<n;i++)
	{
		for(int j=0;j<n;j++)
		{
			pre_density[i][j]=0;
		}
	}
}

void Fluid2D::SwapD()
{
	float ** tmp;
	tmp=density; density=pre_density; pre_density=tmp;
}

void Fluid2D::SwapU()
{
	float ** tmp;
	tmp=u; u=pre_u; pre_u=tmp;
}

void Fluid2D::SwapV()
{
	float ** tmp;
	tmp=v; v=pre_v; pre_v=tmp;
}
